package com.zhang.flink.kafka;

import com.zhang.flink.example.source.ClickSource;
import org.apache.flink.api.common.typeutils.base.StringSerializer;
import org.apache.flink.api.common.typeutils.base.VoidSerializer;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.TwoPhaseCommitSinkFunction;
import java.io.BufferedWriter;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;

/**
 * @title:
 * @author: zhang
 * @date: 2022/3/30 23:11
 */
// 自己实现两阶段提交
// FlinkKafkaProducer的实现思路
public class My2PCFlinkKafkaProducer {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        env.enableCheckpointing(5000L);

        env
                .addSource(new ClickSource())
                .map(r -> r.name)
                .addSink(new TransactionalFileSink());

        env.execute();
    }

    public static class TransactionalFileSink extends TwoPhaseCommitSinkFunction<String, String, Void> {

        //缓冲区
        private BufferedWriter transactionWriter;

        public TransactionalFileSink() {
            super(StringSerializer.INSTANCE, VoidSerializer.INSTANCE);
        }

        //数据流的第一条数据到达的触发调用
        //每次事务commit完之后，触发调用开启新事务
        //创建临时文件
        @Override
        protected String beginTransaction() throws Exception {
            long timeNow = System.currentTimeMillis();
            int taskIdx = this.getRuntimeContext().getIndexOfThisSubtask();
            String transactionFile = timeNow + "-" + taskIdx;
            Path tFilePath = Paths.get("/Users/apple/zhang/filetemp/" + transactionFile);
            Files.createFile(tFilePath);
            this.transactionWriter = Files.newBufferedWriter(tFilePath);
            System.out.println("create tx");
            return transactionFile;
        }

        //预提交
        //将缓冲区中的数据flush到临时文件
        @Override
        protected void preCommit(String transaction) throws Exception {
            transactionWriter.flush();
            transactionWriter.close();
        }

        //回滚机制
        //删除临时文件
        @Override
        protected void abort(String transaction) {
            Path tFilePath = Paths.get("/Users/apple/zhang/filetemp/" + transaction);
            if (Files.exists(tFilePath)) {
                try {
                    Files.delete(tFilePath);
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }


        //当sink算子接收到flink程序检查点完成通知是，正式提交commit
        //降临时文件名改正正式文件名
        @Override
        protected void commit(String transaction) {
            Path tFilePath = Paths.get("/Users/apple/zhang/filetemp/" + transaction);
            if (Files.exists(tFilePath)) {
                try {
                    Path cFilePath = Paths.get("/Users/apple/zhang/filetarget/" + transaction);
                    Files.move(tFilePath, cFilePath);
                    System.out.println("commit complete");
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }

        //每来一条数据调用一次
        //将数据写入缓存区中
        @Override
        protected void invoke(String transaction, String value, Context context) throws Exception {
            transactionWriter.write(value + "\n");
        }


    }
}
